import torch
from collections import deque
import matplotlib.pyplot as plt
import numpy as np

def calculate_distances(tensor):
    """
    计算每个点与其上面点、左边点、左上角点的欧氏距离
    :param tensor: 输入的三维张量，形状为 (1, channel, h, w)
    :return: 上、左、左上角的欧氏距离张量
    """
    tensor = tensor.squeeze(0)
    h, w = tensor.shape[1:]

    # 计算与上面点的欧氏距离
    up_distances = torch.zeros((h, w), dtype=torch.float32)
    up_distances[1:, :] = torch.sqrt(torch.sum((tensor[:, 1:, :] - tensor[:, :-1, :]) ** 2, dim=0))

    # 计算与左边点的欧氏距离
    left_distances = torch.zeros((h, w), dtype=torch.float32)
    left_distances[:, 1:] = torch.sqrt(torch.sum((tensor[:, :, 1:] - tensor[:, :, :-1]) ** 2, dim=0))

    # 计算与左上角点的欧氏距离
    top_left_distances = torch.zeros((h, w), dtype=torch.float32)
    top_left_distances[1:, 1:] = torch.sqrt(torch.sum((tensor[:, 1:, 1:] - tensor[:, :-1, :-1]) ** 2, dim=0))

    return up_distances, left_distances, top_left_distances

def threshold_distances(up_distances, left_distances, top_left_distances, threshold):
    """
    根据阈值标记距离低于阈值的位置为 True
    :param up_distances: 与上面点的欧氏距离张量
    :param left_distances: 与左边点的欧氏距离张量
    :param top_left_distances: 与左上角点的欧氏距离张量
    :param threshold: 距离阈值
    :return: 标记后的上、左、左上角的布尔张量
    """
    up_mask = up_distances < threshold
    left_mask = left_distances < threshold
    top_left_mask = top_left_distances < threshold
    return up_mask, left_mask, top_left_mask

def flood_fill(mask):
    """
    使用洪水填充算法找出相连的 True 区域
    :param mask: 布尔类型的掩码张量
    :return: 标记后的区域张量，每个相连区域有唯一的编号
    """
    h, w = mask.shape
    labels = torch.zeros((h, w), dtype=torch.int32)
    label = 1

    for i in range(h):
        for j in range(w):
            if mask[i, j] and labels[i, j] == 0:
                stack = deque([(i, j)])
                while stack:
                    x, y = stack.pop()
                    if 0 <= x < h and 0 <= y < w and mask[x, y] and labels[x, y] == 0:
                        labels[x, y] = label
                        stack.extend([(x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)])
                label += 1

    return labels

def visualize_and_save(labels, output_path):
    """
    将标记后的区域图可视化并保存为图片
    :param labels: 标记后的区域张量
    :param output_path: 图片保存路径
    """
    # 生成随机颜色映射
    num_labels = int(labels.max()) + 1
    colors = np.random.rand(num_labels, 3)
    colors[0] = [0, 0, 0]  # 背景设为黑色

    # 根据标签索引颜色
    colored_image = colors[labels.numpy()]

    # 保存图片
    plt.imsave(output_path, colored_image)

# 加载 .pt 文件
tensor = torch.load('image_tensor.pt')
tensor=tensor.to('cuda')
print(tensor.shape)
threshold = 1.0

# 计算距离
up_distances, left_distances, top_left_distances = calculate_distances(tensor)
print('get_dist')
# 根据阈值标记
up_mask, left_mask, top_left_mask = threshold_distances(up_distances, left_distances, top_left_distances, threshold)

# 综合三个方向的掩码
combined_mask = up_mask | left_mask | top_left_mask
print('get mask')
# 洪水填充找出相连区域
combined_labels = flood_fill(combined_mask)

# 可视化并保存结果
output_path = 'output_image.png'
visualize_and_save(combined_labels, output_path)
print(f"结果已保存为 {output_path}")
    